#include "FluidSolver.h"
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <memory.h>

CFluidSolver::CFluidSolver(int N, float dt, float diff, float visc)
:m_N(N), m_dt(dt), m_diff(diff), m_visc(visc), size((N + 2) * (N + 2) * (N + 2))
{
	m_dens		= new float[size];
	m_dens_prev	= new float[size];
	m_u			= new float[size];
	m_u_prev	= new float[size];
	m_v			= new float[size];
	m_v_prev	= new float[size];
	m_w			= new float[size];
	m_w_prev	= new float[size];
	curl		= new float[size];
	if (!m_dens || !m_dens_prev || !m_u || !m_u_prev || !m_v || !m_v_prev || !curl)
	{
		fprintf(stderr, "cannot allocate data\n");
	}
	ClearData();

	//build structure of sparse matrix A in A.x = b
	matA = new NRsparseMat(size, size, size * 7);
}

CFluidSolver::~CFluidSolver(void)
{
	FreeMemory();
}

void CFluidSolver::FreeMemory()
{
	if (m_u) delete m_u;
	if (m_u_prev) delete m_u_prev;
	if (m_v) delete m_v;
	if (m_v_prev) delete m_v_prev;
	if (m_w) delete m_w;
	if (m_w_prev) delete m_w_prev;
	if (m_dens) delete m_dens;
	if (m_dens_prev) delete m_dens_prev;
	if (curl) delete curl;

	//delete sparse matrix A
	if (matA) delete matA;
}

void CFluidSolver::ClearData()
{
	for (int i = 0; i < size; i++)
	{
		m_u[i] = m_u_prev[i] = m_v[i] = m_v_prev[i] = m_w[i] = m_w_prev[i] = 0.f;
		m_dens[i] = m_dens_prev[i] = curl[i] = 0.f;
	}
}

void CFluidSolver::LinearSolver(int b, float * x, float * x0, float a, float c)
{
	int i, j, k, l;
	for (l = 0; l < 20; l++)
	{
		for (i = 1; i <= m_N; i++)
		{
			for (j = 1; j <= m_N; j++)
			{
				for (k = 1; k <= m_N; k++)
				{
					x[IX3(i, j, k)] = (x0[IX3(i, j, k)] + a * (
						x[IX3(i - 1, j, k)] + x[IX3(i + 1, j, k)] + 
						x[IX3(i, j - 1, k)] + x[IX3(i, j + 1, k)] +
						x[IX3(i, j, k - 1)] + x[IX3(i, j, k + 1)])) / c;

				}
			}
		}
		SetBoundary(b, x);
	}
}

int DI[] = {-1, 0 , 0 , 0, 0, 0, 1};
int DJ[] = {0 , -1, 0 , 0, 0, 1, 0};
int DK[] = {0 , 0 , -1, 0, 1, 0, 0};

void CFluidSolver::PCBCGSolver(int b, float *x, float *x0, float a, float c)
{
	//build matrix A in A.root = B
	int u, v, i, j, k;
	int t1 = m_N + 2;
	int t2 = (m_N + 2) * (m_N + 2);
	int nval = 0;
	for (u = 0; u < size; u++)
	{
		matA->col_ptr[u] = nval;
		i = u / t2;
		j = (u % t2) / t1;
		k = (u % t2) % t1;
		for (v = 0; v < 7; v++)
		{
			int di = i + DI[v];
			int dj = j + DJ[v];
			int dk = k + DK[v];
			int r = IX3(di,dj,dk);
			if (r >= 0 && r < size)
			{
				if (r == u)
					matA->val[nval] = a;
				else matA->val[nval] = c;
				matA->row_ind[nval] = r;
				nval++;
			}
		}
	}
	matA->col_ptr[size] = nval;

	//build b(a.k.a x0)
	VecDoub B(size, 0.);
	for (i = 0; i < size; i++)
		B[i] = x0[i];
	//build root
	VecDoub_IO root(size, 0.);
	for (i = 0; i < size; i++)
		root[i] = x0[i];
	//solve
	Int iter;
	Doub err;
	NRsparseLinbcg linbcg(*matA);
	linbcg.solve(B, root, 1, 1.0e-1, size, iter, err);

	//copy to x
	for (i = 0; i < size; i++)
	{
		x[i] = (float)root[i];
	}

	//set boundary
	SetBoundary(b, x);
}

void CFluidSolver::VelocitySolver()
{
	//AddSource(m_u, m_u_prev); AddSource(m_v, m_v_prev); AddSource(m_w, m_w_prev);
	//add in buoyancy force
	Buoyancy(m_v_prev, m_dens);
	AddSource(m_v, m_v_prev);

	//add in vorticity confinement force
	VorticityConfinement(m_u_prev, m_v_prev, m_w_prev);
	AddSource(m_u, m_u_prev); AddSource(m_v, m_v_prev); AddSource(m_w, m_w_prev);

	//diffuse
	SWAP(m_u, m_u_prev); Diffuse(1, m_u, m_u_prev, m_visc);
	SWAP(m_v, m_v_prev); Diffuse(2, m_v, m_v_prev, m_visc);
	SWAP(m_w, m_w_prev); Diffuse(3, m_w, m_w_prev, m_visc);
	Project(m_u, m_v, m_w, m_u_prev, m_v_prev);

	//avect
	SWAP(m_u, m_u_prev); SWAP(m_v, m_v_prev); SWAP(m_w, m_w_prev);
	Advect(1, m_u, m_u_prev, m_u_prev, m_v_prev, m_w_prev); 
	Advect(2, m_v, m_v_prev, m_u_prev, m_v_prev, m_w_prev);
	Advect(3, m_w, m_w_prev, m_u_prev, m_v_prev, m_w_prev);
	Project(m_u, m_v, m_w, m_u_prev, m_v_prev);

	//fill m_v_prev with 0
	memset(m_v_prev, 0, size * sizeof(float));
}
void CFluidSolver::DensitySolver()
{
	AddSource(m_dens, m_dens_prev);
	//diffuse
	SWAP(m_dens, m_dens_prev); Diffuse(0, m_dens, m_dens_prev, m_diff);
	//avect
	SWAP(m_dens, m_dens_prev); Advect(0, m_dens, m_dens_prev, m_u, m_v, m_w);

	//fill m_dens_prev with 0
	memset(m_dens_prev, 0, size * sizeof(float));
}

void CFluidSolver::AddSource(float * x, float * x0)
{
	for (int i = 0; i < size; i++)
	{
		x[i] += m_dt * x0[i];
	}
}

void CFluidSolver::Buoyancy(float * Fbuoy, float * temp)
{
	float Tamb = 0;
	float a = 0.000625f;
	//float b = 0.025f;
	float b = 0.5f;

	// sum all temperatures
	int i, j, k;
	for (i = 1; i <= m_N; i++)
	{
		for (j = 1; j <= m_N; j++)
		{
			for (k = 1; k <= m_N; k++)
			{
				Tamb += temp[IX3(i, j, k)];
			}
		}
	}

	// get average temperature
	Tamb /= (m_N * m_N * m_N);

	// for each cell compute buoyancy force
	for (i = 1; i <= m_N; i++)
	{
		for (j = 1; j <= m_N; j++)
		{
			for (k = 1; k <= m_N; k++)
			{
				Fbuoy[IX3(i, j, k)] = -a * temp[IX3(i, j, k)] + b * (temp[IX3(i, j, k)] - Tamb);
			}
		}
	}
}

void CFluidSolver::VorticityConfinement(float * Fvc_x, float * Fvc_y, float * Fvc_z)
{
	float dc_dx, dc_dy, dc_dz;
	float length;
	float x, y, z;
	float Ncx, Ncy, Ncz;

	// Calculate magnitude of curl(u,v,w) for each cell. (|w|)
	int i, j, k, ijk;
	for (i = 1; i <= m_N; i++)
	{
		for (j = 1; j <= m_N; j++)
		{
			for (k = 1; k <= m_N; k++)
			{
				ijk = IX3(i, j, k);
				//curlx = dw/dy - dv/dz
				x = Fvc_x[ijk] = (m_w[IX3(i, j + 1, k)] - m_w[IX3(i, j - 1, k)]) * 0.5f -
					(m_v[IX3(i, j, k + 1)] - m_v[IX3(i, j, k - 1)]) * 0.5f;
				//curly = du/dz - dw/dx;
				y = Fvc_y[ijk] = (m_u[IX3(i, j, k + 1)] - m_u[IX3(i, j , k - 1)]) * 0.5f -
					(m_w[IX3(i + 1, j, k)] - m_w[IX3(i - 1, j , k)]) * 0.5f;
				//curlz = dv/dx - du/dy;
				z = Fvc_z[ijk] = (m_v[IX3(i + 1, j, k)] - m_v[IX3(i - 1, j, k)]) * 0.5f -
					(m_u[IX3(i, j + 1, k)] - m_u[IX3(i, j - 1, k)]) * 0.5f;
				curl[ijk] = sqrt(x * x + y * y + z * z);
			}
		}
	}

	for (i = 2; i < m_N; i++)
	{
		for (j = 2; j < m_N; j++)
		{
			for (k = 2; k < m_N; k++)
			{

				// Find derivative of the magnitude (n = del |c|)
				dc_dx = (curl[IX3(i + 1, j, k)] - curl[IX3(i - 1, j, k)]) * 0.5f;
				dc_dy = (curl[IX3(i, j + 1, k)] - curl[IX3(i, j - 1, k)]) * 0.5f;
				dc_dz = (curl[IX3(i, j, k + 1)] - curl[IX3(i, j, k - 1)]) * 0.5f;

				// Calculate vector length. (|n|)
				// Add small factor to prevent divide by zeros.
				length = (float)sqrt(dc_dx * dc_dx + dc_dy * dc_dy + dc_dz * dc_dz)
					+ 0.000001f;

				// N = ( n/|n| )
				dc_dx /= length;
				dc_dy /= length;
				dc_dz /= length;

				// N x c
				Ncx = dc_dy * Fvc_z[ijk] - dc_dz * Fvc_y[ijk];
				Ncy = dc_dz * Fvc_x[ijk] - dc_dx * Fvc_z[ijk];
				Ncz = dc_dx * Fvc_y[ijk] - dc_dy * Fvc_x[ijk];
				Fvc_x[ijk] = Ncx;
				Fvc_y[ijk] = Ncy;
				Fvc_z[ijk] = Ncz;
			}
		}
	}
}

void CFluidSolver::Diffuse(int b, float * c, float * c0, float diff)
{
	float a = m_dt * diff * m_N * m_N * m_N;
	LinearSolver(b, c, c0, a, 1 + 6 * a);

	//float cc = -m_dt * diff * m_N * m_N * m_N;
	//float aa = 1 + 6 * cc;
	//PCBCGSolver(b, c, c0, aa, cc);
}

void CFluidSolver::Advect(int b, float * d, float * d0, float * u, float * v, float * w)
{
	int i, j, k, i0, j0, k0, i1, j1, k1;
	float xx, yy, zz, dt0;
	float sx0, sx1, sy0, sy1, sz0, sz1, v0, v1;

	dt0 = m_dt * m_N;
	for (i = 1; i <= m_N; i++)
	{
		for (j = 1; j <= m_N; j++)
		{
			for (k = 1; k <= m_N; k++)
			{
				xx = i - dt0 * u[IX3(i, j, k)];
				yy = j - dt0 * v[IX3(i, j, k)];
				zz = k - dt0 * w[IX3(i, j, k)];
				if (xx < 0.5f) xx = 0.5f; if (xx > m_N + 0.5f) xx = m_N + 0.5f; i0 = (int)xx; i1 = i0 + 1;
				if (yy < 0.5f) yy = 0.5f; if (yy > m_N + 0.5f) yy = m_N + 0.5f; j0 = (int)yy; j1 = j0 + 1;
				if (zz < 0.5f) zz = 0.5f; if (zz > m_N + 0.5f) zz = m_N + 0.5f; k0 = (int)zz; k1 = k0 + 1;

				sx1 = xx - i0; sx0 = 1 - sx1;
				sy1 = yy - j0; sy0 = 1 - sy1;
				sz1 = zz - k0; sz0 = 1 - sz1;

				v0 = sx0 * (sy0 * d0[IX3(i0,j0,k0)] + sy1 * d0[IX3(i0,j1,k0)]) + sx1 * (sy0 * d0[IX3(i1,j0,k0)] + sy1 * d0[IX3(i1,j1,k0)]);
				v1 = sx0 * (sy0 * d0[IX3(i0,j0,k1)] + sy1 * d0[IX3(i0,j1,k1)]) + sx1 * (sy0 * d0[IX3(i1,j0,k1)] + sy1 * d0[IX3(i1,j1,k1)]);
				d[IX3(i,j,k)] = sz0 * v0 + sz1 * v1;
			}
		}
	}
	SetBoundary(b, d);
}

/*void CFluidSolver::BFECC(int b, float * d, float * d0, float * u, float * v)
{
	int i, j, i0, j0, i1, j1;
	float x, y, s0, t0, s1, t1, dt0;

	//float f_phi, b_phi;

	dt0 = m_dt * m_N;

	//test upwinding advection
	//float gra;
	//for (i = 1; i <= m_N; i++)
	//{
	//	for (j = 1; j <= m_N; j++)
	//	{
	//		gra = 0.5f * (d0[IX(i, j + 1)] - d0[IX(i, j - 1)]) * v[IX(i, j)] 
	//			+ 0.5f * (d0[IX(i + 1, j)] - d0[IX(i - 1, j)]) * u[IX(i, j)];
	//		d[IX(i, j)] = d0[IX(i, j)] - dt0 * gra;
	//	}
	//}	

	//u, v, phi, phi~
	for (i = 1; i <= m_N; i++)
	{
		for (j = 1; j <= m_N; j++)
		{
			x = i - dt0 * u[IX(i, j)]; y = j - dt0 * v[IX(i, j)];
			if (x < 0.5f) x = 0.5f; if (x > m_N + 0.5f) x = m_N + 0.5f; i0 = (int)x; i1 = i0 + 1;
			if (y < 0.5f) y = 0.5f; if (y > m_N + 0.5f) y = m_N + 0.5f; j0 = (int)y; j1 = j0 + 1;
			s1 = x - i0; s0 = 1 - s1; t1 = y - j0; t0 = 1 - t1;
			d[IX(i, j)] = s0 * (t0 * d0[IX(i0, j0)] + t1 * d0[IX(i0, j1)])+
				s1 * (t0 * d0[IX(i1, j0)] + t1 * d0[IX(i1, j1)]);
		}
	}

	////-u, -v, phi~, phi_
	//for (i = 1; i <= m_N; i++)
	//{
	//	for (j = 1; j <= m_N; j++)
	//	{
	//		x = i - dt0 * -u[IX(i, j)]; y = j - dt0 * -v[IX(i, j)];
	//		if (x < 0.5f) x = 0.5f; if (x > m_N + 0.5f) x = m_N + 0.5f; i0 = (int)x; i1 = i0 + 1;
	//		if (y < 0.5f) y = 0.5f; if (y > m_N + 0.5f) y = m_N + 0.5f; j0 = (int)y; j1 = j0 + 1;
	//		s1 = x - i0; s0 = 1 - s1; t1 = y - j0; t0 = 1 - t1;
	//		d0[IX(i, j)] = (3 * d0[IX(i, j)] - (s0 * (t0 * d[IX(i0, j0)] + t1 * d[IX(i0, j1)])+
	//			s1 * (t0 * d[IX(i1, j0)] + t1 * d[IX(i1, j1)]))) / 2.0;
	//	}
	//}

	////phi(n + 1)
	//for (i = 1; i <= m_N; i++)
	//{
	//	for (j = 1; j <= m_N; j++)
	//	{
	//		x = i - dt0 * u[IX(i, j)]; y = j - dt0 * v[IX(i, j)];
	//		if (x < 0.5f) x = 0.5f; if (x > m_N + 0.5f) x = m_N + 0.5f; i0 = (int)x; i1 = i0 + 1;
	//		if (y < 0.5f) y = 0.5f; if (y > m_N + 0.5f) y = m_N + 0.5f; j0 = (int)y; j1 = j0 + 1;
	//		s1 = x - i0; s0 = 1 - s1; t1 = y - j0; t0 = 1 - t1;
	//		d[IX(i, j)] = s0 * (t0 * d0[IX(i0, j0)] + t1 * d0[IX(i0, j1)])+
	//			s1 * (t0 * d0[IX(i1, j0)] + t1 * d0[IX(i1, j1)]);
	//	}
	//}
	
	SetBoundary(b, d );

}*/

void CFluidSolver::Project(float * u, float * v, float * w, float * p, float * div)
{
	int i, j, k;
	float h = 1.0f / m_N;
	for (i = 1; i <= m_N; i++)
	{
		for (j = 1; j <= m_N; j++)
		{
			for (k = 1; k <= m_N; k++)
			{
				//different from smoke64 ?????
				div[IX3(i, j, k)] = -0.5f * h * (
					u[IX3(i + 1, j, k)] - u[IX3(i - 1 ,j, k)] + 
					v[IX3(i, j + 1, k)] - v[IX3(i, j - 1, k)] +
					w[IX3(i, j, k + 1)] - w[IX3(i, j, k - 1)]);
				p[IX3(i, j, k)] = 0;
			}
		}
	}
	SetBoundary(0, div); SetBoundary(0, p);

	LinearSolver(0, p, div, 1, 6);
	//PCBCGSolver(0, p, div, -6, 1);

	for (i = 1; i <= m_N; i++)
	{
		for (j = 1; j <= m_N; j++)
		{
			for (k = 1; k <= m_N; k++)
			{
				//different from smoke64 ?????
				u[IX3(i, j, k)] -= 0.5f * m_N * (p[IX3(i + 1, j, k)] - p[IX3(i - 1, j, k)]);
				v[IX3(i, j, k)] -= 0.5f * m_N * (p[IX3(i, j + 1, k)] - p[IX3(i, j - 1, k)]);
				w[IX3(i, j, k)] -= 0.5f * m_N * (p[IX3(i, j, k + 1)] - p[IX3(i, j, k - 1)]);
			}
		}
	}
	SetBoundary(1, u); SetBoundary(2, v); SetBoundary(3, w);
}

void CFluidSolver::SetBoundary(int b, float * x)
{
	int i, j;
	for (i = 1; i <= m_N; i++)
	{
		for (j = 1; j <= m_N; j++)
		{
			x[IX3(0, i, j)]			= b == 1 ? -x[IX3(1, i, j)] : x[IX3(1, i, j)];
			x[IX3(m_N + 1, i, j)]	= b == 1 ? -x[IX3(m_N, i, j)] : x[IX3(m_N, i, j)];
			x[IX3(i, 0, j)]			= b == 2 ? -x[IX3(i, 1, j)] : x[IX3(i, 1, j)];
			x[IX3(i, m_N + 1, j)]	= b == 2 ? -x[IX3(i, m_N, j)] : x[IX3(i, m_N, j)];
			x[IX3(i, j, 0)]			= b == 3 ? -x[IX3(i, j, 1)] : x[IX3(i, j, 1)];
			x[IX3(i, j, m_N + 1)]	= b == 3 ? -x[IX3(i, j, m_N)] : x[IX3(i, j, m_N)];
			//x[IX3(0, i, j)]			= 0;
			//x[IX3(m_N + 1, i, j)]	= 0;
			//x[IX3(i, 0, j)]			= 0;
			//x[IX3(i, m_N + 1, j)]	= 0;
			//x[IX3(i, j, 0)]			= b == 3 ? -x[IX3(i, j, 1)] : x[IX3(i, j, 1)];
			//x[IX3(i, j, m_N + 1)]	= 0;
		}
	}
	x[IX3(0,0,0)]				= (x[IX3(1,0,0)]			+x[IX3(0,1,0)]				+x[IX3(0,0,1)])				/3;
	x[IX3(0,m_N+1,0)]			= (x[IX3(1,m_N+1,0)]		+x[IX3(0,m_N,0)]			+x[IX3(0,m_N+1,1)])			/3;
	x[IX3(m_N+1,0,0)]			= (x[IX3(m_N,0,0)]			+x[IX3(m_N+1,1,0)]			+x[IX3(m_N+1,0,1)])			/3;
	x[IX3(m_N+1,m_N+1,0)]		= (x[IX3(m_N,m_N+1,0)]		+x[IX3(m_N+1,m_N,0)]		+x[IX3(m_N+1,m_N+1,1)])		/3;
	x[IX3(0,0,m_N+1)]			= (x[IX3(1,0,m_N+1)]		+x[IX3(0,1,m_N+1)]			+x[IX3(0,0,m_N)])			/3;
	x[IX3(0,m_N+1,m_N+1)]		= (x[IX3(1,m_N+1,m_N+1)]	+x[IX3(0,m_N,m_N+1)]		+x[IX3(0,m_N+1,m_N)])		/3;
	x[IX3(m_N+1,0,m_N+1)]		= (x[IX3(m_N,0,m_N+1)]		+x[IX3(m_N+1,1,m_N+1)]		+x[IX3(m_N+1,0,m_N)])		/3;
	x[IX3(m_N+1,m_N+1,m_N+1)]	= (x[IX3(m_N,m_N+1,m_N+1)]	+x[IX3(m_N+1,m_N,m_N+1)]	+x[IX3(m_N+1,m_N+1,m_N)])	/3;
}

void CFluidSolver::SetSource(int i, int j, int k, float val)
{
	m_dens_prev[IX3(i,j,k)] = val;
}